/**
 * 
 */
package util;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkNotNull;

import java.lang.reflect.Array;
import java.util.Collection;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

import com.google.common.collect.AbstractIterator;

/**
 * @author Joshua Lockerman
 * @param <V> 
 *
 */
public class TrieSet<V> implements Set<Map.Entry<String, V>> {
    
    protected Root root;
    private int size;
    
    public TrieSet()
    {
        root = new Root();
        size = 0;
    }
    
    @Override
    public boolean add(Entry<String, V> e)
    {
        TrieSearchResult location = root.findSuffixTree(e.getKey());
        if(location != null) {
            if( !(location.finishedSearchingKey() && location.finishedSearchingSuffix())) {
                location.addToNode(e);
                size++;
                return true;
            }
        }
        return false;
    }
    
    @Override
    public boolean addAll(Collection<? extends Entry<String, V>> c)
    {
        boolean added = true;
        for(Entry<String, V> entry : c) {
            added |= add(entry);
        }
        return added;
    }
    
    @Override
    public void clear()
    {
        root = new Root();
        size = 0;
    }
    
    @Override
    public boolean contains(Object o)
    {
        if( !(o instanceof TrieSet.MapPair)) {
            return false;
        }
        MapPair pair = (TrieSet.MapPair) o;
        TrieSearchResult location = root.findSuffixTree(pair.getKey());
        if(location != null && location.finishedSearchingKey()
                && location.finishedSearchingSuffix()) {
            return location.getNode().getEntry().equals(o);
        }
        return false;
    }
    
    @Override
    public boolean containsAll(Collection<?> c)
    {
        boolean containsAll = true;
        for(Object o : c) {
            containsAll &= contains(o);
        }
        return containsAll;
    }
    
    public TrieSearchResult findSuffixTree(String key)
    {
        return root.findSuffixTree(key);
    }
    
    public Root getRoot()
    {
        return root;
    }
    
    @Override
    public boolean isEmpty()
    {
        return size == 0;
    }
    
    @Override
    public Iterator iterator()
    {
        return new Iterator();
    }
    
    public V put(String key, V value)
    {
        checkArgument( !key.isEmpty());
        checkNotNull(value);
        
        TrieSearchResult location = root.findSuffixTree(key);
        checkNotNull(location);
        size++;
        return location.addToNode(new MapPair(key, value));
    }
    
    @Override
    public boolean remove(Object o)
    {
        if( !(o instanceof Entry)) {
            return false;
        }
        Entry entry = (Entry) o;
        if(root.remove(entry) != null) {
            size--;
            return true;
        }
        return false;
    }
    
    @Override
    public boolean removeAll(Collection<?> c)
    {
        boolean changed = false;
        for(Object o : c) {
            changed |= remove(o);
        }
        return changed;
    }
    
    public V removeEntry(String s)
    {
        TrieSearchResult r = root.findSuffixTree(s);
        if(r.finishedSearchingKey() && r.finishedSearchingSuffix()) {
            return r.getNode().removeData();
        }
        return null;
    }
    
    @Override
    public boolean retainAll(Collection<?> c)
    {
        Root newRoot = new Root();
        int newSize = 0;
        for(Object o : c) {
            if(contains(o)) {
                V oldVal = newRoot.add((Entry<String, V>) o);
                if(oldVal == null) {
                    newSize++;
                }
            }
        }
        boolean changed = newSize != size();
        root = newRoot;
        size = newSize;
        return changed;
    }
    
    @Override
    public int size()
    {
        return size;
    }
    
    @Override
    public Object[] toArray()
    {
        Object[] ar = new Object[size()];
        Iterator it = iterator();
        int index = 0;
        while(it.hasNext()) {
            ar[index] = it.next();
            index++;
        }
        return ar;
    }
    
    @Override
    public <T> T[] toArray(T[] a)
    {
        if(a.length < size()) {
            Class<?> type = a.getClass().getComponentType();
            a = (T[]) Array.newInstance(type, size());
        }
        Iterator it = iterator();
        int index = 0;
        while(it.hasNext()) {
            a[index] = (T) it.next();
            index++;
        }
        while(index < a.length) {
            a[index] = null;
            index++;
        }
        return a;
    }
    
    public class Iterator extends AbstractIterator<Map.Entry<String, V>> implements
    java.util.Iterator<Entry<String, V>> {
        
        private final LinkedList<java.util.Iterator<Entry<Character, Node>>> stack;
        
        public Iterator()
        {
            stack = new LinkedList<java.util.Iterator<Entry<Character, Node>>>();
            stack.push(root.childIterator());
        }
        
        @Override
        protected Entry<String, V> computeNext()
        {
            if( !stack.isEmpty()) {
                java.util.Iterator<Entry<Character, Node>> it = stack.peek();
                while( !it.hasNext() && stack.size() > 1) {
                    stack.pop();
                    it = stack.peek();
                }
                if(it.hasNext()) {
                    Node node = it.next().getValue();
                    stack.push(node.childIterator());
                    if(node.getData() != null) {
                        Entry<String, V> next = node.getEntry();
                        return next;
                    }
                    else {
                        return computeNext();
                    }
                }
            }
            return endOfData();
        }
        
    }
    
    /**
     * A key-value pair.
     * 
     * @author Joshua Lockerman
     * 
     */
    public class MapPair implements Map.Entry<String, V>,
    Comparable<Map.Entry<String, ? extends V>> {
        
        private String key;
        private V value;
        
        /**
         * Sole constructor.
         * 
         * @param key
         * @param value
         */
        public MapPair(String key, V value)
        {
            this.key = key;
            this.value = value;
        }
        
        @Override
        public int compareTo(java.util.Map.Entry<String, ? extends V> o)
        {
            return key.compareTo(o.getKey());
        }
        
        @Override
        public boolean equals(Object obj)
        {
            if(obj == this) {
                return true;
            }
            
            if(obj instanceof TrieSet.MapPair) {
                MapPair pair = (TrieSet.MapPair) obj;
                if(key.equals(pair.getKey())) {
                    if(value != null) {
                        if(pair.value != null) {
                            return value.getClass().equals(pair.getValue().getClass());
                        }
                    }
                    else if(pair.value == null) {
                        return true;
                    }
                }
            }
            
            return false;
        }
        
        @Override
        public String getKey()
        {
            return key;
        }
        
        @Override
        public V getValue()
        {
            return value;
        }
        
        @Override
        public int hashCode()
        {
            return key.hashCode();
        }
        
        /**
         * Sets the key of the key-value pair to be s
         * 
         * @param s
         *            the new key.
         */
        public void setKey(String s)
        {
            key = s;
        }
        
        @Override
        public V setValue(V arg0)
        {
            V old = value;
            value = arg0;
            return old;
        }
    }
    
    /**
     * Object used as the result from the find method in Node.
     * It contains the last node searched, the index of the 
     * last character of the key that was checked and the index of the 
     * last character of the suffix that was checked 
     * 
     * @author Joshua Lockerman
     */
    public class TrieSearchResult {
        private final int keySearchLength;
        private final int suffixSearchLength;
        private final Node node;
        private final String key;
        
        private TrieSearchResult(int keyindex, int suffixIndex, Node node, String key)
        {
            checkNotNull(node);
            checkNotNull(key);
            checkArgument(keyindex >= 0);
            checkArgument(suffixIndex >= 0);
            
            this.keySearchLength = keyindex;
            this.suffixSearchLength = suffixIndex;
            this.node = node;
            this.key = key;
        }
        
        /**
         * Continues the search by appending a given suffix to the search key.<\br>
         * ie: has the same result as searching for <code>{@link #getKey()} + suffix<code>
         * 
         * @note changes to the underlying trie could cause undefined results
         * @param suffix the suffix to be added to the search key.
         * @return a TrieSearchResult representing the results of the search
         */
        public final TrieSearchResult appendToSearch(String suffix)
        {
            if( !finishedSearchingKey() && finishedSearchingSuffix()) {
                return new TrieSearchResult(keySearchLength, suffixSearchLength, node, key + suffix);
            }
            return node.findSuffixTree(key + suffix, keySearchLength, keySearchLength);
        }
        
        public final boolean finishedSearchingKey()
        {
            return keySearchLength >= key.length();
        }
        
        public final boolean finishedSearchingSuffix()
        {
            return suffixSearchLength >= node.getSuffix().length();
        }
        
        public final String getKey()
        {
            return key;
        }
        
        /**
         * The number of characters of the key matched against the Trie 
         * @return the number of characters of the prefix of the key has in common
         * with a path through the Trie 
         */
        public final int getKeySearchLength()
        {
            return keySearchLength;
        }
        
        /**
         * Returns the last node searched.
         * @return the last node searched.
         */
        public Node getNode()
        {
            return node;
        }
        
        /**
         * The number of characters the partial suffix contained in the {@link Node} from
         *  {@link TrieSearchResult#getNode() getNode()} has in common with the searched key.
         * @return the number of characters of the suffix that were matched with the key
         */
        public final int getSuffixSearchLength()
        {
            return suffixSearchLength;
        }
        
        private V addToNode(Map.Entry<String, V> entry)
        {
            checkArgument(entry.getKey().equals(key));
            return node.put(this, entry);
        }
    }
    
    /**
     * Node for TreeTrieMap.
     * 
     * @author Joshua Lockerman
     * 
     */
    public class Node {
        
        private Node parent;
        private char nodeKey;
        private LinkedHashMap<Character, Node> children;
        private Entry<String, V> entry;
        private String suffix;
        
        /**
         * Constructor for the root node.
         * 
         */
        private Node()
        {
            this.parent = null;
            children = new LinkedHashMap<Character, Node>();
            entry = new MapPair("", null);
            suffix = "";
        }
        
        private Node(Node node, char nodeKey, Entry<String, V> entry, String suffix,
                LinkedHashMap<Character, Node> children)
        {
            this.parent = node;
            this.nodeKey = nodeKey;
            this.entry = entry;
            this.suffix = suffix;
            this.children = children;
            java.util.Iterator<Entry<Character, Node>> it = childIterator();
            while(it.hasNext()) {
                Entry<Character, Node> child = it.next();
                child.getValue().setParent(this);
            }
        }
        
        /**
         * Constructor for Node.
         * 
         * @param parent
         *            the parent of this node.
         * @param key
         *            the key for the value stored in this node.
         * @param value
         *            the value stored in this node.
         * @param suffixStart
         *            the index that the suffix stored in this node should start at. Must be <= key.length()
         */
        protected Node(Node parent, char nodeKey, TrieSearchResult location, Entry<String, V> value)
        {
            if(parent == null) {
                throw new NullPointerException("Attempted to set a null parent in TrieMap.Node");
            }
            checkNotNull(value);
            
            this.parent = parent;
            this.nodeKey = nodeKey;
            children = new LinkedHashMap<Character, Node>();
            entry = value;
            suffix = location.getKey().substring(location.getKeySearchLength() + 1);
        }
        
        public java.util.Iterator<Entry<Character, Node>> childIterator()
        {
            return children.entrySet().iterator();
        }
        
        public V getData()
        {
            return entry.getValue();
        }
        
        private Entry<String, V> getEntry()
        {
            return entry;
        }
        
        public Node getParent()
        {
            return parent;
        }
        
        public String getSuffix()
        {
            return suffix;
        }
        
        public boolean hasChildren()
        {
            return children.size() > 0;
        }
        
        private void addChild(TrieSearchResult location, Map.Entry<String, V> value)
        {
            char key = location.getKey().charAt(location.getKeySearchLength());
            children.put(key, new Node(this, key, location, value));
        }
        
        private char getNodeKeyChar()
        {
            String key = entry.getKey();
            return key.charAt(key.length() - suffix.length() - 1);
        }
        
        private V put(TrieSearchResult location, Map.Entry<String, V> value)
        {
            if( !location.finishedSearchingSuffix()) {
                split(location, value);
            }
            else if( !location.finishedSearchingKey()) {
                addChild(location, value);
            }
            else {
                V old = getData();
                this.setData(value.getValue());
                return old;
            }
            return null;
        }
        
        private void setData(V data)
        {
            entry.setValue(data);
        }
        
        private void setParent(Node parent)
        {
            this.parent = parent;
        }
        
        private void split(TrieSearchResult location, Entry<String, V> value)
        {
            String key = location.getKey();
            int keySearchLength = location.getKeySearchLength();
            int suffixSearchLength = location.getSuffixSearchLength();
            // Add the old value from this node into a new child
            char childKey = suffix.charAt(suffixSearchLength);
            Node child = new Node(this, childKey, getEntry(), suffix.substring(suffixSearchLength + 1),children);
            children = new LinkedHashMap<Character, Node>();
            children.put(childKey, child);
            suffix = suffix.substring(0, suffixSearchLength);
            //FIXME
            if( !location.finishedSearchingKey()) {
                entry = new MapPair(key.substring(0, keySearchLength), null);
                // Create a new child for the new key.
                Node newChild = new Node(this, key.charAt(keySearchLength), location, value);
                children.put(key.charAt(keySearchLength), newChild);
            }
            else {
                // add the key-value pair to this node.
                this.entry = value;
            }
        }
        
        protected final TrieSearchResult findSuffixTree(String key, int keyIndex)
        {
            int length = Math.min(key.length() - keyIndex, suffix.length());
            int offset = 0;
            boolean equal = true;
            while(offset < length && equal) {
                equal &= (key.charAt(keyIndex + offset) == suffix.charAt(offset));
                if(equal) {
                    offset++;
                }
            }
            if(equal && keyIndex + offset != key.length()) {
                // if we have searched the entire suffix but not the entire key
                Node node = children.get(key.charAt(keyIndex + offset));
                if(node != null) {
                    return node.findSuffixTree(key, keyIndex + offset + 1);
                }
            }
            
            return new TrieSearchResult(keyIndex + offset, offset, this, key);
        }
        
        protected final TrieSearchResult findSuffixTree(String key, int keyIndex, int suffixIndex)
        {
            int length = Math.min(key.length() - keyIndex, suffix.length() - suffixIndex);
            int offset = 0;
            boolean equal = true;
            while(offset < length && equal) {
                equal &= (key.charAt(keyIndex + offset) == suffix.charAt(suffixIndex + offset));
                if(equal) {
                    offset++;
                }
            }
            if(equal && keyIndex + offset != key.length()) {
                // if we have searched the entire suffix but not the entire key
                Node node = children.get(key.charAt(keyIndex + offset));
                if(node != null) {
                    return node.findSuffixTree(key, keyIndex + offset + 1);
                }
            }
            
            return new TrieSearchResult(keyIndex + offset, offset, this, key);
        }
        
        protected V removeData()
        {
            V oldValue = null;
            if(getData() != null) {
                oldValue = getData();
                setData(null);
                if(children.isEmpty()) {
                    //if this has no data and no children it's useless
                    getParent().children.remove(nodeKey);
                }
                else if(children.size() == 1) {
                    // only one child so we can merge the child with this node
                    Entry<Character, Node> ent = childIterator().next();
                    this.entry = ent.getValue().entry;
                    suffix += ent.getKey() + ent.getValue().suffix;
                    children = ent.getValue().children;
                    java.util.Iterator<Entry<Character, Node>> it = childIterator();
                    while(it.hasNext()) {
                        Entry<Character, Node> child = it.next();
                        child.getValue().setParent(this);
                    }
                }
            }
            return oldValue;
        }
        
        public boolean isLeaf()
        {
            return children.isEmpty();
        }
    }
    
    protected class Root extends Node {
        
        public Root()
        {
            super();
        }
        
        public V add(Map.Entry<String, V> entry)
        {
            checkNotNull(entry.getValue());
            checkNotNull(entry.getKey());
            checkArgument( !entry.getKey().isEmpty());
            
            TrieSearchResult n = findSuffixTree(entry.getKey());
            return n.addToNode(entry);
        }
        
        public V get(String key)
        {
            checkNotNull(key);
            checkArgument( !key.isEmpty());
            TrieSearchResult result = findSuffixTree(key);
            if(result.finishedSearchingKey() && result.finishedSearchingSuffix()) {
                return result.getNode().getData();
            }
            return null;
        }
        
        public TrieSearchResult findSuffixTree(String key)
        {
            return findSuffixTree(key, 0);
        }
        
        public V remove(Map.Entry<String, V> entry)
        {
            checkNotNull(entry.getValue());
            checkArgument( !entry.getKey().isEmpty());
            
            TrieSearchResult n = findSuffixTree(entry.getKey());
            if(n != null && n.finishedSearchingKey() && n.finishedSearchingSuffix()) {
                if(n.getNode().getData() != null) {
                    V v = n.getNode().getData();
                    n.getNode().removeData();
                    return v;
                }
            }
            
            return null;
        }
        
        @Override
        protected V removeData()
        {
            return null;
        }
    }
}
